package joe;

import java.lang.foreign.Arena;
import java.lang.foreign.FunctionDescriptor;
import java.lang.foreign.Linker;
import java.lang.foreign.SymbolLookup;
import java.lang.invoke.MethodHandle;
import org.jetbrains.annotations.NotNull;
import static java.lang.foreign.ValueLayout.JAVA_INT;
import static java.lang.foreign.ValueLayout.JAVA_LONG;

// ffmpeg -i input_any.mp4 -c:a libopus -b:a 48k -c:v libsvtav1 -crf 40 -g 300 output_av1.mp4
// ffmpeg -i input_av1.mp4 -c copy -map_metadata -1 output.avif
// ffmpeg -i input_any.mp4 -c:a pcm_s16le -ac 2 -ar 48000 output.wav
// opaenc 48000 input.wav output.opa
// ref: https://trac.ffmpeg.org/wiki/Encode/AV1#SVT-AV1
// ref: https://github.com/BtbN/FFmpeg-Builds/releases
public class ResourceVideo implements Resource {
	private static final @NotNull MethodHandle mhAvifOpenFile;
	private static final @NotNull MethodHandle mhAvifDecodeFile;
	private static final @NotNull MethodHandle mhAvifCloseFile;

	static {
		var libAvif = SymbolLookup.libraryLookup(System.mapLibraryName("avif"), Arena.global());
		var linker = Linker.nativeLinker();
		// avifDecoder* avifOpenFile(const char* filename, uint64_t offset, uint64_t size)
		mhAvifOpenFile = linker.downcallHandle(libAvif.find("avifOpenFile").orElseThrow(),
				FunctionDescriptor.of(JAVA_LONG, JAVA_LONG, JAVA_LONG, JAVA_LONG));
		// int avifDecodeFile(avifDecoder* decoder, int frameIdx)
		mhAvifDecodeFile = linker.downcallHandle(libAvif.find("avifDecodeFile").orElseThrow(),
				FunctionDescriptor.of(JAVA_INT, JAVA_LONG, JAVA_INT));
		// void avifCloseFile(avifDecoder* decoder)
		mhAvifCloseFile = linker.downcallHandle(libAvif.find("avifCloseFile").orElseThrow(),
				FunctionDescriptor.ofVoid(JAVA_LONG));
	}

	private final long decoder;
	private final @NotNull Runnable cleaner;

	public ResourceVideo(@NotNull String filename, long offset, long size) {
		var bytes = filename.getBytes();
		if (bytes.length >= Engine.TMP_SIZE)
			throw new IllegalArgumentException("too long filename");
		var ptr = Engine.current().getTmpBuf().ptr();
		UnsafeUtil.copyMemory(bytes, ptr);
		UnsafeUtil.putByte(ptr + bytes.length, (byte)0);
		long decoder;
		try {
			decoder = (long)mhAvifOpenFile.invokeExact(ptr, offset, size);
		} catch (Throwable e) { // rethrow
			throw UnsafeUtil.forceThrow(e);
		}
		if (decoder <= 0)
			throw new IllegalStateException("avifOpenFile failed: " + decoder);
		this.decoder = decoder;
		cleaner = () -> {
			try {
				mhAvifCloseFile.invokeExact(decoder);
			} catch (Throwable e) { // rethrow
				UnsafeUtil.forceThrow(e);
			}
		};
	}

	@Override
	public @NotNull Runnable getCleaner() {
		return cleaner;
	}

/*
struct avifDecoder {
+00 avifCodecChoice codecChoice;
+04 int maxThreads;
+08 avifDecoderSource requestedSource;
+0C avifBool allowProgressive;
+10 avifBool allowIncremental;
+14 avifBool ignoreExif;
+18 avifBool ignoreXMP;
+1C uint32_t imageSizeLimit;
+20 uint32_t imageDimensionLimit;
+24 uint32_t imageCountLimit;
+28 avifStrictFlags strictFlags;
+30 avifImage* image;
+38 int imageIndex;
+3C int imageCount;
+40 avifProgressiveState progressiveState;
+48 avifImageTiming imageTiming;
+48   uint64_t timescale;
+50   double pts;
+58   uint64_t ptsInTimescales;
+60   double duration;
+68   uint64_t durationInTimescales;
+70 uint64_t timescale;
+78 double duration;
+80 uint64_t durationInTimescales;
+88 int repetitionCount;
+8C avifBool alphaPresent;
+90 // ...
}
*/

	public int getCurFrameIdx() {
		return UnsafeUtil.getInt(decoder + 0x38);
	}

	public int getFrameCount() {
		return UnsafeUtil.getInt(decoder + 0x3c);
	}

	public double getCurFrameTime() {
		return UnsafeUtil.getDouble(decoder + 0x50);
	}

	public double getCurFrameDuration() {
		return UnsafeUtil.getDouble(decoder + 0x60);
	}

	public double getDuration() {
		return UnsafeUtil.getDouble(decoder + 0x78);
	}

	public boolean hasAlpha() {
		return UnsafeUtil.getInt(decoder + 0x8c) != 0;
	}

	/**
	 * @return pointers in the Image are only available before next decoding or closing
	 */
	public @NotNull Image decodeNextImage() {
		return decodeImage(-1);
	}

	/**
	 * @return pointers in the Image are only available before next decoding or closing
	 */
	public @NotNull Image decodeImage(int frameIdx) {
		int r;
		try {
			r = (int)mhAvifDecodeFile.invokeExact(decoder, frameIdx);
		} catch (Throwable e) { // rethrow
			throw UnsafeUtil.forceThrow(e);
		}
		if (r != 0)
			throw new IllegalStateException("avifDecodeFile failed: " + r);
		var p = UnsafeUtil.getLong(decoder + 0x30);
		if (p == 0)
			throw new IllegalStateException("avifDecodeFile failed: no image");
		return new Image(p);
	}

	/*
	struct avifImage {
	+00 uint32_t width;
	+04 uint32_t height;
	+08 uint32_t depth;
	+0C avifPixelFormat yuvFormat;
	+10 avifRange yuvRange;
	+14 avifChromaSamplePosition yuvChromaSamplePosition;
	+18 uint8_t* yuvPlanes[AVIF_PLANE_COUNT_YUV=3];
	+30 uint32_t yuvRowBytes[AVIF_PLANE_COUNT_YUV=3];
	+3C avifBool imageOwnsYUVPlanes;
	+40 uint8_t* alphaPlane;
	+48 uint32_t alphaRowBytes;
	+4C avifBool imageOwnsAlphaPlane;
	+50 avifBool alphaPreMultiplied;
	+58 avifRWData icc;
	+58   uint8_t* data;
    +60   size_t size;
    +68 avifColorPrimaries colorPrimaries;
    +6A avifTransferCharacteristics transferCharacteristics;
    +6C avifMatrixCoefficients matrixCoefficients;
	+6E // ...
	}
	*/
	public static final class Image {
		public static final int YUV_FORMAT_444 = 1;
		public static final int YUV_FORMAT_422 = 2;
		public static final int YUV_FORMAT_420 = 3;
		public static final int YUV_FORMAT_400 = 4;
		public static final int SAMPLE_POS_UNKNOWN = 0;
		public static final int SAMPLE_POS_VERTICAL = 1;
		public static final int SAMPLE_POS_COLOCATED = 2;
		public static final int MAT_COEFF_IDENTITY = 0;    // YUV=RGB
		public static final int MAT_COEFF_BT709 = 1;       // BT.709
		public static final int MAT_COEFF_UNSPECIFIED = 2; // default
		public static final int MAT_COEFF_FCC = 4;         // FCC USFC 73.682
		public static final int MAT_COEFF_BT470BG = 5;     // BT.470-6 System BG
		public static final int MAT_COEFF_BT601 = 6;       // BT.601
		public static final int MAT_COEFF_SMPTE240 = 7;    // SMPTE ST 240
		public static final int MAT_COEFF_BT2020_NCL = 9;  // BT.2020 (non-constant luminance)

		public final int width;
		public final int height;
		public final int depth; // 8/10/12
		public final int yuvFormat; // enum
		public final boolean yuvFullRange; // full:[0,2^depth-1]; rangeLimited:biasY=16<<(depth-8);rangeY=219<<(depth-8);rangeUV=224<<(depth-8)
		public final int samplePos; // enum
		public final long yData, uData, vData, aData; // pointers
		public final int yStride, uStride, vStride, aStride; // bytes
		public final boolean alphaPreMultiplied;
		public final short matrixCoefficients; // 0:identity 1:bt709 2:unspecified 4:fcc 5:bt470bg 6:bt601 7:smpte240 9:bt2020ncl

		private Image(long p) {
			width = UnsafeUtil.getInt(p);
			height = UnsafeUtil.getInt(p + 0x4);
			depth = UnsafeUtil.getInt(p + 0x8);
			yuvFormat = UnsafeUtil.getInt(p + 0xc);
			yuvFullRange = UnsafeUtil.getInt(p + 0x10) == 1;
			samplePos = UnsafeUtil.getInt(p + 0x14);
			yData = UnsafeUtil.getLong(p + 0x18);
			uData = UnsafeUtil.getLong(p + 0x20);
			vData = UnsafeUtil.getLong(p + 0x28);
			yStride = UnsafeUtil.getInt(p + 0x30);
			uStride = UnsafeUtil.getInt(p + 0x34);
			vStride = UnsafeUtil.getInt(p + 0x38);
			aData = UnsafeUtil.getLong(p + 0x40);
			aStride = UnsafeUtil.getInt(p + 0x48);
			alphaPreMultiplied = UnsafeUtil.getInt(p + 0x50) != 0;
			matrixCoefficients = UnsafeUtil.getShort(p + 0x6c);
		}

		public @NotNull Mat4 createCoeffMat() {
			return createCoeffMat(matrixCoefficients, yuvFullRange);
		}

		// Y  = (iY  - biasY ) / rangeY
		// UV = (iUV - biasUV) / rangeUV (biasUV=1<<(depth-1))
		// R = Y + V * (2-kr*2);
		// G = Y - V * (kr*2-kr*kr*2)/kg + U * (kb*2-kb*kb*2)/kg; (kg=1-kr-kb)
		// B = Y                         + U * (2-kb*2);
		public static @NotNull Mat4 createCoeffMat(int matrixCoefficients, boolean yuvFullRange) {
			final var mat = new Mat4();
			final float kr, kb;
			switch (matrixCoefficients) {
			//@formatter:off
			case MAT_COEFF_IDENTITY:   return mat.identity();
			case MAT_COEFF_BT709:      kr=0.2126f; kb=0.0722f; break;
			case MAT_COEFF_FCC:        kr=0.3000f; kb=0.1100f; break;
			case MAT_COEFF_SMPTE240:   kr=0.2120f; kb=0.0870f; break;
			case MAT_COEFF_BT2020_NCL: kr=0.2627f; kb=0.0593f; break;
			default:                   kr=0.2990f; kb=0.1140f; break;
			//@formatter:on
			}
			final var biasUV = -128f / 255;
			final var kg2 = 2 / (1 - kr - kb);
			final var v11 = (kb * kb - kb) * kg2;
			final var v12 = 2 - kb * 2;
			final var v20 = 2 - kr * 2;
			final var v21 = (kr * kr - kr) * kg2;
			final var v30 = biasUV * v20;
			final var v31 = biasUV * (v11 + v21);
			final var v32 = biasUV * v12;
			if (yuvFullRange) {
				return mat
						.setRow(0, 1, 1, 1, 0)
						.setRow(1, 0, v11, v12, 0)
						.setRow(2, v20, v21, 0, 0)
						.setRow(3, v30, v31, v32, 1);
			}
			final var rangeY = 255f / 219;
			final var rangeUV = 255f / 224;
			final var biasY = -16f / 219;
			//noinspection SuspiciousNameCombination
			return mat
					.setRow(0, rangeY, rangeY, rangeY, 0)
					.setRow(1, 0, v11 * rangeUV, v12 * rangeUV, 0)
					.setRow(2, v20 * rangeUV, v21 * rangeUV, 0, 0)
					.setRow(3, v30 + biasY, v31 + biasY, v32 + biasY, 1);
		}
	}
}
